#!/usr/bin/env python3
# Copyright (c) 2019-2020 Windy Albert


import hashlib
import binascii
import hmac
import ecdsa
import os
import base58check
import bech32

from keyTree import *


class addrBTC(object):
    def __init__(self, obj):
        self.obj = obj

    def getWIF(self,pa):
        xprv = keyTree().getChildpriKey(self.obj,pa)
        xprv_hex = base58check.b58decode(xprv["xprv"])
        prikey_hex = xprv_hex[46:78].hex()
        data = "80"+prikey_hex + "01" #compressed
        sha256_1 = hashlib.new('sha256',bytes.fromhex(data))
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        wif_b58 = base58check.b58encode(bytes.fromhex(data+(sha256_2.hexdigest())[0:8]))

        path = self.obj["path"] + "/" + pa
        wif = str(wif_b58,'utf-8')
        
        return {"path":path,"wif":wif}

    def getPubkeyID(self,pubkey):
        # wif_hex = base58check.b58decode(wif)
        # prikey = wif_hex[1:33]
        # sk = ecdsa.SigningKey.from_string( bytes.fromhex(prikey.hex()), curve=ecdsa.SECP256k1)    
        # vk = sk.get_verifying_key() #this is your verification key (public key)
        # public_key = vk.to_string("compressed").hex()
        sha256  = hashlib.new('sha256',bytes.fromhex(pubkey))
        hash160 = hashlib.new('ripemd160',bytes.fromhex(sha256.hexdigest()))
        return (hash160.hexdigest())

    def getPubkey(self,wif):
        wif_hex = base58check.b58decode(wif)
        prikey = wif_hex[1:33]
        sk = ecdsa.SigningKey.from_string( bytes.fromhex(prikey.hex()), curve=ecdsa.SECP256k1)    
        vk = sk.get_verifying_key() #this is your verification key (public key)
        public_key = vk.to_string("compressed").hex()
        return public_key

class addrBTC_p2pkh(addrBTC):
    def getWallet(self,path):
        wallet = super().getWIF(path)
        pubkey = super().getPubkey(wallet["wif"])
        wallet["addr"] = self.getAddr( pubkey )
        return wallet

    def getAddr(self,pubkey):
        pubkeyID = super().getPubkeyID(pubkey)
        data = "00" + pubkeyID
        sha256_1 = hashlib.new('sha256',bytes.fromhex(data))
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        addr = base58check.b58encode(bytes.fromhex(data+(sha256_2.hexdigest())[0:8]))
        return addr.decode('utf-8')

class addrBTC_p2wpkh(addrBTC):
    def getWallet(self,path):
        wallet = super().getWIF(path)
        pubkey = super().getPubkey(wallet["wif"])
        wallet["addr"] = self.getAddr( pubkey )
        return wallet

    def getAddr(self,pubkey):
        pubkeyID = super().getPubkeyID(pubkey)
        data = bech32.convertbits(bytes.fromhex(pubkeyID),8,5)
        data.insert(0,0)
        addr = bech32.bech32_encode("bc",data)
        return addr

class addrBTC_p2wpkh_p2sh(addrBTC):
    def getWallet(self,path):
        wallet = super().getWIF(path)
        pubkey = super().getPubkey(wallet["wif"])
        wallet["addr"] = self.getAddr( pubkey )
        return wallet

    def getAddr(self,pubkey):
        pubkeyID = super().getPubkeyID(pubkey)
        sha256  = hashlib.new('sha256',bytes.fromhex("0014" + pubkeyID))
        scriptId = "05" + hashlib.new('ripemd160',bytes.fromhex(sha256.hexdigest())).hexdigest()
        sha256_1 = hashlib.new('sha256',bytes.fromhex(scriptId))
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        addr = base58check.b58encode(bytes.fromhex(scriptId+(sha256_2.hexdigest())[0:8]))
        return addr.decode('utf-8')

class addrBTC_p2wsh(addrBTC):
    def getWallet(self,path):
        wallet = super().getWIF(path)
        wallet["addr"] = self.getAddr( super().getPubkey(wallet["wif"]) )
        return wallet

    def getAddr(self,pubkey):
        mnScript = hashlib.new('sha256',bytes.fromhex("5121"+pubkey+"51ae")).hexdigest()        
        data = bech32.convertbits(bytes.fromhex(mnScript),8,5)
        data.insert(0,0)
        addr = bech32.bech32_encode("bc",data)
        return addr

class addrBTC_p2wsh_p2sh(addrBTC):
    def getWallet(self,path):
        wallet = super().getWIF(path)
        wallet["addr"] = self.getAddr( super().getPubkey(wallet["wif"]) )
        return wallet

    def getAddr(self,pubkey):
        script = "0020"+hashlib.new('sha256',bytes.fromhex("5121" + pubkey+"51ae")).hexdigest()
        sha256  = hashlib.new('sha256',bytes.fromhex(script)).hexdigest()
        scriptId = hashlib.new('ripemd160',bytes.fromhex(sha256)).hexdigest()
        data = "05" + scriptId # add version
        sha256_1 = hashlib.new('sha256',bytes.fromhex(data))
        sha256_2 = hashlib.new('sha256',bytes.fromhex(sha256_1.hexdigest()))
        addr = base58check.b58encode(bytes.fromhex(data+(sha256_2.hexdigest())[0:8]))
        return addr.decode('utf-8')


def main():
    print("This is demo ...")

if __name__ == '__main__':
    main()